#include "stdafx.h"
#include "cppflibinc.h"
#include "TestMiscFilePath.h"

using namespace cppflib;
using namespace cppflib::os;
using namespace cppflib::collections;
using namespace cppflib::exceptions;
using namespace cppflib::net;
using namespace cppflib::net::ssl;
using namespace cppflib::threading;
using namespace cppflib::util;
using namespace cppflib::io;

#define SERVERPORT          6229 
#define SERVERADDR          "127.0.0.1" 
#define CLIENTADDR          "127.0.0.1"

class MyClientPrivateKeyCallback : public ISSLPrivateKeyPwdCallback
{
public:
    virtual void GetPwd(CFString &pwd)
    {
       pwd = "zxcv1234";  // fill in password for private key
    }
};

class MyServerPrivateKeyCallback : public ISSLPrivateKeyPwdCallback
{
public:
    virtual void GetPwd(CFString &pwd)
    {
       pwd = "zxcv1234";  // fill in password for private key
    }
};

static void PrintCertExtensions(security::cert::X509Certificate & x509Cert)
{
   debug::Print(debug::DL_DEBUG, _S("extension count = %d\n"), x509Cert.GetExtensionCount());

   for (int i = 0; i < x509Cert.GetExtensionCount(); ++i) {
        CFString name;
        if (x509Cert.GetExtensionName(i, name))
           debug::Print(debug::DL_DEBUG, _S("Extension name[%d] = %s\n"), i, static_cast<pcwstr_t>(name));

        security::cert::X509Certificate::extensiontype_t type = security::cert::X509Certificate::EXT_UNKNOWN;
        if (!x509Cert.GetExtensionType(i, type))
           continue;

        if (type == security::cert::X509Certificate::EXT_LIST) {

           NameValueList extValList;
           if (x509Cert.GetExtensionValueList(i, extValList)) {
              for (int j = 0; j < static_cast<int>(extValList.GetCount()); ++j) {
                 debug::Print(debug::DL_DEBUG, _S("   Extmval[%d] = %s:%s\n"), j, static_cast<pcwstr_t>(*extValList.GetKey(j)),
                                                                                 static_cast<pcwstr_t>(*extValList.GetVal(j)));
              }
           }

        }
        else if (type == security::cert::X509Certificate::EXT_STRING) {

           CFString val;
           if (x509Cert.GetExtensionValueString(i, val)) 
              debug::Print(debug::DL_DEBUG, _S("   ExtString = %s\n"), static_cast<pcwstr_t>(val));

        }
        else if (type == security::cert::X509Certificate::EXT_RAW) {
           BinBuf binBuf;

           if (x509Cert.GetExtensionValueRaw(i, binBuf)) {
              CFString s;
              u8_t * pBuf = binBuf.GetRawBuf();
              u8_t * pBufEnd = binBuf.GetValidEnd();
              while (pBuf < pBufEnd) 
                 s.AppendFormat(_S("%02x "), *pBuf++);
              debug::Print(debug::DL_DEBUG, _S("   Extraw = %s\n"), static_cast<pcwstr_t>(s));
           }

        }
   }
}


class MySSLPeerCertCallback : public ISSLPeerCertCallback
{
private:
    bool isPrintCert;
    int certCount;

public:
    MySSLPeerCertCallback()
    {
       isPrintCert = false;
       certCount = 0;
    }

    ~MySSLPeerCertCallback()
    {
       debug::Print(_S("Clean cert file created...\n"));
       if (certCount > 0) {
          int num = certCount;
          while (num > 0) {
             --num;

             CFString name;
             GetCertName(name, num);
             debug::Print(debug::DL_DEBUG, _S("Deleting %s ...\n"), static_cast<pcwstr_t>(name));
             os::DelFile(static_cast<pcwstr_t>(name));
          }
       }
    }

    explicit MySSLPeerCertCallback(bool isPrint)
    {
       isPrintCert = isPrint;
    }

    void GetCertName(CFString &name, int num)
    {
       name.Format(_S("cert_%d.cer"), num);
    }

    void SaveCert(security::cert::X509Certificate &x509Cert)
    {
       ArrayByte *pBuf = x509Cert.GetEncoded();

       if (pBuf == NULL) {
          debug::Print(_S("ooops! no encoded cert.\n"));
          return;
       }

       CFString name;
       GetCertName(name, certCount);
       debug::Print(debug::DL_DEBUG, _S("Saving %s ...\n"), static_cast<pcwstr_t>(name));
       FileStream *f = new FileStream(name.ToString(), FileStream::FA_WRITE);
       f->SetOutputEncoding(FileStream::ENC_NONE);
       f->Write(*pBuf);
       f->Close();

       ++certCount;
    }

    void PrintPeerCert(bool isCertOK, security::cert::X509Certificate &x509Cert)
    {
       CFString tempName;

       debug::Print(_S("--------------------------------------------------\n"));
       debug::Print(debug::DL_DEBUG, _S("cert OK = %s\n"), isCertOK ? _S("true") : _S("false"));
       debug::Print(debug::DL_DEBUG, _S("version = %d\n"), x509Cert.GetVersion());
       debug::Print(debug::DL_DEBUG, _S("issuer = %s\n"), static_cast<pcwstr_t>(x509Cert.GetIssuerName()));
       if (x509Cert.GetOneValueFromIssuerName(CFString("emailAddress"), tempName))
           debug::Print(debug::DL_DEBUG, _S("issuer emailAddress = %s\n"), static_cast<pcwstr_t>(tempName));
       debug::Print(debug::DL_DEBUG, _S("subject = %s\n"), static_cast<pcwstr_t>(x509Cert.GetSubjectName()));
       if (x509Cert.GetOneValueFromSubjectName(CFString("CN"), tempName))
           debug::Print(debug::DL_DEBUG, _S("subject CN = %s\n"), static_cast<pcwstr_t>(tempName));
       CFString notBefore(x509Cert.GetNotBefore().ToString());
       debug::Print(debug::DL_DEBUG, _S("not before = %s\n"), static_cast<pcwstr_t>(notBefore));
       CFString notAfter(x509Cert.GetNotAfter().ToString());
       debug::Print(debug::DL_DEBUG, _S("not after = %s\n"), static_cast<pcwstr_t>(notAfter));
       debug::Print(debug::DL_DEBUG, _S("cert valid = %s\n"), x509Cert.IsValid(*DateTime::UtcNow()) ? _S("true") : _S("false"));
       PrintCertExtensions(x509Cert);
    }

    virtual void ReadPeerCert(bool isCertOK, security::cert::X509Certificate &x509Cert)
    {
       if (!isPrintCert)
          return;

       PrintPeerCert(isCertOK, x509Cert);

       SaveCert(x509Cert);
    }

    /**
     *  Called when connection is established with peer
     */
    virtual void ReadPeerCertAfterConnection(security::cert::X509Certificate &x509Cert)
    {
       
    }
    
};


class MyServerSSLPeerCertCallback : public ISSLPeerCertCallback
{
private:
    bool isPrintCert;

public:
    MyServerSSLPeerCertCallback()
    {
       isPrintCert = false;
    }

    explicit MyServerSSLPeerCertCallback(bool isPrintCert)
    {
       this->isPrintCert = isPrintCert;
    }

    virtual void ReadPeerCert(bool isCertOK, security::cert::X509Certificate &x509Cert)
    {
       if (!isPrintCert)
          return;

       debug::Print(_S("--------------------------------------------------\n"));
       debug::Print(debug::DL_DEBUG, _S("client issuer = %s\n"), static_cast<pcwstr_t>(x509Cert.GetIssuerName()));
       debug::Print(debug::DL_DEBUG, _S("client subject = %s\n"), static_cast<pcwstr_t>(x509Cert.GetSubjectName()));
    }

    /**
     *  Called when connection is established with peer
     */
    virtual void ReadPeerCertAfterConnection(security::cert::X509Certificate &x509Cert)
    {
       if (!isPrintCert)
          return;

       debug::Print(_S("****--------------------------------------------------****\n"));
       debug::Print(debug::DL_DEBUG, _S("+++ client subject = %s\n"), static_cast<pcwstr_t>(x509Cert.GetSubjectName()));
    }
};

static char tempStr[] = "1234qwerasdf1234\n";

static void TestSSLClient(bool isTrustServer)
{
   
   SSLSocket *pSock = NULL;
   try {
      SSLContext *sslCtx = NULL;

      if (isTrustServer) {
         sslCtx = new SSLContext();
      }
      else {
         sslCtx = new SSLContext(false, CFString(ROOTCERT), false);
         sslCtx->SetCertChain(CFString(CLIENTCERT), new MyClientPrivateKeyCallback());
      }

      pSock = sslCtx->CreateClientSocket(NULL, new MySSLPeerCertCallback(isTrustServer));
      pSock->Create(CFString(CLIENTADDR), 0);
      //debug::Print(debug::DL_INFO, _S("Connecting ...\n"));
      pSock->Connect(CFString(SERVERADDR), SERVERPORT, 3000);
      //debug::Print(debug::DL_INFO, _S("Connect OK ...\n"));
      Thread::Sleep(1000);
      pSock->Send(reinterpret_cast<const u8_t*>(tempStr), CFString::Strlen(tempStr), 0);
   }
   catch(SSLSocketException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }
   catch(SSLException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }
   
   if (pSock) {
      pSock->Close();
   }
   
}

static void __SendDataFunc(Event *pEvent, bool isTrustServer)
{
   cppmempool::MemPoolWrap localPool;

   if (isTrustServer)
      debug::Print(_S("Always trust SSL server test ...\n"));
   debug::Print(debug::DL_DEBUG, _S("client thread -- 0x%x\n"), Thread::GetCurrentThreadId());

   // wait until server ready
   pEvent->Wait();

   Thread::Sleep(100);

   // send data to server
   TestSSLClient(isTrustServer);
}

static int SendDataThread(ptr_t p)
{
   __SendDataFunc(reinterpret_cast<Event*>(p), false);
   return 0;
}

static int SendDataThreadTrustServer(ptr_t p)
{
   __SendDataFunc(reinterpret_cast<Event*>(p), true);
   return 0;
}


static int ServerThread(ptr_t p)
{
   cppmempool::MemPoolWrap localPool;

   debug::Print(debug::DL_DEBUG, _S("server thread -- 0x%x\n"), Thread::GetCurrentThreadId());

   SSLSocket *pSock = reinterpret_cast<SSLSocket*>(p);
   cppmempool::AddToCurPool(pSock);

   try {
      const int BUFSIZE = 128;
      u8_t buf[BUFSIZE];
            
      Thread::Sleep(1500);

      pSock->WaitHandshake(5 * 1000);
      int nrRead = pSock->ReceiveTimeout(3000, buf, BUFSIZE, 0);
      //printf("nrRead = %d\n", nrRead);
      assert(nrRead > 0);
      buf[nrRead] = 0;
      CFString a(reinterpret_cast<const char*>(buf));
      CFString b(reinterpret_cast<const char*>(tempStr));
      assert(a == b);
   }
   catch(SSLSocketException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }
   catch(SSLException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }

   pSock->Close();
   return 0;
}

/**
 *  Spawn 3 client threads to connect to this server. 
 *  Server will then spawn a thread for each incoming client.
 */
static void TestSSLServerClient()
{
   cppmempool::MemPoolWrap localPool;

   debug::Print(debug::DL_DEBUG, _S("OpenSSL version = %s\n"), static_cast<pcwstr_t>(SSLContext::GetOpenSSLVersion()));

   // start client thread to send data
   Thread a((ThreadEntryPoint(&SendDataThread)));
   Event *startSendEvent = new Event();  // a event
   a.Start(reinterpret_cast<ptr_t>(startSendEvent));

   Thread b((ThreadEntryPoint(&SendDataThread)));
   Event *startSendEvent2 = new Event();  // a event
   b.Start(reinterpret_cast<ptr_t>(startSendEvent2));

   Thread c((ThreadEntryPoint(&SendDataThread)));
   Event *startSendEvent3 = new Event();  // a event
   c.Start(reinterpret_cast<ptr_t>(startSendEvent3)); 
 

   const int TH_COUNT = 3;
   Thread *sTh[TH_COUNT];

   try {
      SSLContext *sslCtx = new SSLContext(true, CFString(ROOTCERT), true);  // require client authentication
      sslCtx->SetCertChain(CFString(SERVERCERT), new MyServerPrivateKeyCallback());

      SSLServerSocket *pServSock = sslCtx->CreateServerSocket(new MyServerSSLPeerCertCallback());
      pServSock->Create(CFString(SERVERADDR), SERVERPORT);
      assert(pServSock->IsValidSock());
      assert(pServSock->Listen(5));
      
      // signal it is ready to receive data
      startSendEvent->Set();
      startSendEvent2->Set();
      startSendEvent3->Set();

      // spawn a thread to serve each incoming client
      for (int i = 0; i < TH_COUNT; ++i) {
         SSLSocket *pSock = static_cast<SSLSocket*>(pServSock->Accept());
         sTh[i] = new Thread(ThreadEntryPoint(&ServerThread));
         cppmempool::UnlinkFromPool(pSock);
         sTh[i]->Start(reinterpret_cast<ptr_t>(pSock));
      }

   }
   catch(SSLSocketException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }
   catch(SSLException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }


   for (int i = 0; i < TH_COUNT; ++i) {
      sTh[i]->Join();
   }

   startSendEvent->Close();
   startSendEvent2->Close();
   startSendEvent3->Close();
   
   a.Join();
   b.Join();
   c.Join();

}


static void TestAlwaysTrustServer(void)
{
   cppmempool::MemPoolWrap localPool;

   // start client thread to send data
   Thread a((ThreadEntryPoint(&SendDataThreadTrustServer)));
   Event *startSendEvent = new Event();  // a event
   a.Start(reinterpret_cast<ptr_t>(startSendEvent));

   Thread *pTh = NULL;

   try {
      SSLContext *sslCtx = new SSLContext(true, CFString(ROOTCERT), false);  // doesn't require client authentication
      sslCtx->SetCertChain(CFString(SERVERCERT), new MyServerPrivateKeyCallback());

      SSLServerSocket *pServSock = sslCtx->CreateServerSocket(new MyServerSSLPeerCertCallback());
      pServSock->Create(CFString(SERVERADDR), SERVERPORT);
      assert(pServSock->IsValidSock());
      assert(pServSock->Listen(5));
      
      // signal it is ready to receive data
      startSendEvent->Set();
      
      // spawn a thread to serve each incoming client
      SSLSocket *pSock = static_cast<SSLSocket*>(pServSock->Accept());
      ServerThread(reinterpret_cast<ptr_t>(pSock));
/*
//      pTh = new Thread(&ServerThread);
//      cppmempool::UnlinkFromPool(pSock);
//      pTh->Start(reinterpret_cast<ptr_t>(pSock));
*/      
   }
   catch(SSLSocketException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }
   catch(SSLException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }


   if (pTh)
     pTh->Join();
   startSendEvent->Close();
   a.Join();
   
}

static void TestReadSelfSignCert()
{
   debug::Print(_S("Testing reading X509 cert ...\n"));

   cppmempool::MemPoolWrap localPool;

   debug::Print(_S("******* Reading X509 cert in DER format *******\n"));
   security::cert::X509Certificate * x509Cert = security::cert::X509Certificate::FromDER(CFString(SELFSIGNCERT));
   PrintCertExtensions(*x509Cert);

   debug::Print(_S("******* Reading X509 cert in PEM format *******\n"));
   x509Cert = security::cert::X509Certificate::FromPEM(CFString(SELFSIGNCERTPEM));
   PrintCertExtensions(*x509Cert);
}

/*
 * true:  client -> server
 * false: server -> client
 */
static bool IS_CLIENT_SEND = true;

static void __DataSendRecv(SSLSocket *pSock, bool isClient)
{
   const int SEND_COUNT = 10;
   
   if ((IS_CLIENT_SEND && isClient) || (!IS_CLIENT_SEND && !isClient)) {

      Thread::Sleep(200);
      
      // send data to peer
      for (int i = 0; i < SEND_COUNT; ++i) {
         pSock->Send(reinterpret_cast<const u8_t*>(tempStr), CFString::Strlen(tempStr), 0);
         Thread::Sleep(50);
         if (i == 4) {
            assert(pSock->Renegotiate(true, 8 * 1000, true));  // renegotiate in the middle
         }
      }
   }
   else {
      const int BUFSIZE = 128;
      u8_t buf[BUFSIZE];
      for (int i = 0; i < SEND_COUNT; ++i) {
         //int nrRead = pSock->ReceiveTimeout(1000, buf, BUFSIZE, 0);
         int nrRead = pSock->Receive(buf, BUFSIZE, 0);
         assert(nrRead > 0);
         buf[nrRead] = 0;
         CFString a(reinterpret_cast<const char*>(buf));
         CFString b(reinterpret_cast<const char*>(tempStr));
         assert(a == b);
/*
         if (i == 4) {
            Thread::Sleep(10 * 1000);
         }
*/
      }
   }
      
}

static int SendDataClient(ptr_t p)
{
   cppmempool::MemPoolWrap localPool;

   Event *pEvent = reinterpret_cast<Event*>(p);
   
   // wait until server ready
   pEvent->Wait();

   Thread::Sleep(100);

   SSLSocket *pSock = NULL;
   SSLContext *sslCtx = NULL;
   try {
      sslCtx = new SSLContext(false, CFString(ROOTCERT), false);
      sslCtx->SetCertChain(CFString(CLIENTCERT), new MyClientPrivateKeyCallback());
      
      pSock = sslCtx->CreateClientSocket(NULL, NULL);
      pSock->Create(CFString(CLIENTADDR), 0);
      pSock->Connect(CFString(SERVERADDR), SERVERPORT, 3000);      

      //printf("client start\n");

      
      __DataSendRecv(pSock, true);

   }
   catch(SSLSocketException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }
   catch(SSLException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }
   
   if (pSock) {
      pSock->Close();
   }

   return 0;
}



static void TestRenegotiate(int round)
{
   debug::Print(debug::DL_INFO, _S("Testing SSL Renegotiation, round = %d...\n"), round);

   cppmempool::MemPoolWrap localPool;
   // start client thread to send data

   Thread a((ThreadEntryPoint(&SendDataClient)));
   Event *startSendEvent = new Event();  // a event
   a.Start(reinterpret_cast<ptr_t>(startSendEvent));

   SSLContext *sslCtx = NULL;
   try {
      sslCtx = new SSLContext(true, CFString(ROOTCERT), false);  // doesn't require client authentication
      sslCtx->SetCertChain(CFString(SERVERCERT), new MyServerPrivateKeyCallback());

      SSLServerSocket *pServSock = sslCtx->CreateServerSocket(new MyServerSSLPeerCertCallback());
      pServSock->Create(CFString(SERVERADDR), SERVERPORT);
      assert(pServSock->IsValidSock());
      assert(pServSock->Listen(5));
      
      // signal it is ready to receive data
      startSendEvent->Set();
      
      SSLSocket *pSock = static_cast<SSLSocket*>(pServSock->Accept());
      pSock->WaitHandshake(5 * 1000);

      //printf("server start\n");

      __DataSendRecv(pSock, false);

   }
   catch(SSLSocketException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }
   catch(SSLException &exp) {
      debug::Print(debug::DL_DEBUG, _S("Error: %s\n"), exp.GetMsg());
      assert(false);
   }

   startSendEvent->Close();
   a.Join();
   
}

static void TestCheckHostName()
{
   debug::Print(_S("Testing check host name ...\n"));

   assert(IsHostNameMatched(CFString(_S("www.hello.com")), CFString(_S("*.hello.com"))));
   assert(IsHostNameMatched(CFString(_S("www.hello.com")), CFString(_S("w*.hello.com"))));
   assert(IsHostNameMatched(CFString(_S("www.hello.com")), CFString(_S("www.hello.*"))));
   assert(IsHostNameMatched(CFString(_S("www.hello.com.hk")), CFString(_S("www.hello.*"))));
   assert(IsHostNameMatched(CFString(_S("www.hello.com")), CFString(_S("www.hello.com"))));
   assert(! IsHostNameMatched(CFString(_S("www.hello.com.hk")), CFString(_S("*.hello.com"))));
   assert(IsHostNameMatched(CFString(_S("www.hello.com.hk")), CFString(_S("*.hello.com.hk"))));
   assert(! IsHostNameMatched(CFString(_S("www.hello.com")), CFString(_S("abc.hello.com"))));
   assert(! IsHostNameMatched(CFString(_S("www.hello.com")), CFString(_S("*.*.com"))));
   assert(IsHostNameMatched(CFString(_S("www.hello.com")), CFString(_S("www.*.com"))));
}

void TestNetSSL(void)
{
   debug::Print(_S("Testing Net SSL ...\n"));
   
   TestCheckHostName();
   TestReadSelfSignCert();
   TestSSLServerClient();
   TestAlwaysTrustServer();
   for (int i = 0; i < 2; ++i) {
      if (i == 1) 
         IS_CLIENT_SEND = !IS_CLIENT_SEND;
      TestRenegotiate(i);
   }
   
}

